#include "cgzCLua.h"
#include <sqlite3.h>

struct cgzCSqlite
{
    sqlite3 *db;
    std::string m_name;
    std::string m_filename;
};

// 执行非查询SQL语句（CREATE/DROP/ALTER/TRUNCATE等）
static int executeDDL(lua_State *L, sqlite3 *db, const char *sql)
{
    char *errMsg = nullptr;
    int rc = sqlite3_exec(db, sql, nullptr, nullptr, &errMsg);

    if (rc != SQLITE_OK)
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, errMsg);
        sqlite3_free(errMsg);
        return 2;
    }

    lua_pushboolean(L, true);
    return 1;
}

// 执行查询并返回结果
static int executeQuery(lua_State *L, sqlite3 *db, const char *sql)
{
    sqlite3_stmt *stmt;
    int rc = sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);

    if (rc != SQLITE_OK)
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, sqlite3_errmsg(db));
        return 2;
    }

    int num_fields = sqlite3_column_count(stmt);

    // 创建字段名表
    lua_cgz_newa(L, num_fields);
    for (int i = 0; i < num_fields; i++)
    {
        lua_pushstring(L, sqlite3_column_name(stmt, i));
        lua_rawseti(L, -2, i + 1);
    }

    // 创建结果表
    lua_cgz_newt(L);
    int row_index = 1;

    while (sqlite3_step(stmt) == SQLITE_ROW)
    {
        lua_cgz_newa(L, num_fields);

        for (int i = 0; i < num_fields; i++)
        {
            switch (sqlite3_column_type(stmt, i))
            {
            case SQLITE_NULL:
                lua_pushnil(L);
                break;

            case SQLITE_INTEGER:
                lua_pushinteger(L, sqlite3_column_int64(stmt, i));
                break;

            case SQLITE_FLOAT:
                lua_pushnumber(L, sqlite3_column_double(stmt, i));
                break;

            case SQLITE_TEXT:
            {
                const char *text = (const char *)sqlite3_column_text(stmt, i);
                lua_pushstring(L, text);
                break;
            }

            case SQLITE_BLOB:
            {
                const void *blob = sqlite3_column_blob(stmt, i);
                int length = sqlite3_column_bytes(stmt, i);
                lua_pushlstring(L, (const char *)blob, length);
                break;
            }
            }
            lua_rawseti(L, -2, i + 1);
        }

        lua_rawseti(L, -2, row_index++);
    }

    sqlite3_finalize(stmt);
    return 2; // 返回字段名表和结果表
}

// 执行INSERT/UPDATE/DELETE并返回影响行数和最后ID
static int executeDML(lua_State *L, sqlite3 *db, const char *sql)
{
    sqlite3_stmt *stmt = nullptr;
    int rc = sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr);

    if (rc != SQLITE_OK)
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, sqlite3_errmsg(db));
        return 2;
    }

    // 绑定参数（从Lua栈中获取）
    int param_count = sqlite3_bind_parameter_count(stmt);
    for (int i = 1; i <= param_count; i++)
    {
        int pos = i + 2;
        if (lua_isnil(L, pos))
        {
            sqlite3_bind_null(stmt, i);
        }
        else if (lua_isinteger(L, pos))
        {
            sqlite3_bind_int64(stmt, i, lua_tointeger(L, pos));
        }
        else if (lua_isnumber(L, pos))
        {
            sqlite3_bind_double(stmt, i, lua_tonumber(L, pos));
        }
        else if (lua_isstring(L, pos))
        {
            size_t len;
            const char *str = lua_tolstring(L, pos, &len);
            sqlite3_bind_text(stmt, i, str, len, SQLITE_TRANSIENT);
        }
        else if (lua_isuserdata(L, pos))
        {
            // 假设二进制数据通过userdata传递
            void *blob = lua_touserdata(L, pos);
            size_t size = lua_rawlen(L, pos);
            sqlite3_bind_blob(stmt, i, blob, size, SQLITE_TRANSIENT);
        }
        else
        {
            sqlite3_finalize(stmt);
            lua_pushboolean(L, false);
            lua_pushstring(L, "Unsupported parameter type");
            return 2;
        }
    }

    rc = sqlite3_step(stmt);
    if (rc != SQLITE_DONE)
    {
        sqlite3_finalize(stmt);
        lua_pushboolean(L, false);
        lua_pushstring(L, sqlite3_errmsg(db));
        return 2;
    }

    sqlite3_finalize(stmt);

    lua_pushboolean(L, true);
    lua_pushinteger(L, sqlite3_changes(db));
    lua_pushinteger(L, sqlite3_last_insert_rowid(db));
    return 3;
}

static int sqlite_new(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *filename = luaL_checkstring(L, 2);
    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    if (dbs.find(name) != dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection already exists");
        return 2;
    }

    cgzCSqlite *conn = new cgzCSqlite();
    conn->m_name = name;
    conn->m_filename = filename;
    int rc = sqlite3_open(filename, &conn->db);
    if (rc != SQLITE_OK)
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, sqlite3_errmsg(conn->db));
        delete conn;
        return 2;
    }

    dbs[name] = conn;
    lua_pushboolean(L, true);
    return 1;
}

static int sqlite_create(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDDL(L, it->second->db, sql);
}

static int sqlite_drop(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDDL(L, it->second->db, sql);
}

static int sqlite_alter(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDDL(L, it->second->db, sql);
}

static int sqlite_truncate(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDDL(L, it->second->db, sql);
}

static int sqlite_insert(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDML(L, it->second->db, sql);
}

static int sqlite_select(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeQuery(L, it->second->db, sql);
}

static int sqlite_update(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDML(L, it->second->db, sql);
}

static int sqlite_delete(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);
    const char *sql = luaL_checkstring(L, 2);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it == dbs.end())
    {
        lua_pushboolean(L, false);
        lua_pushstring(L, "Database connection not found");
        return 2;
    }

    return executeDML(L, it->second->db, sql);
}

static int sqlite_close(lua_State *L)
{
    const char *name = luaL_checkstring(L, 1);

    cgzCLua *that = cgzCLua::getThat(L);
    auto &dbs = that->sqliteMap;
    auto it = dbs.find(name);
    if (it != dbs.end())
    {
        sqlite3_close(it->second->db);
        delete it->second;
        dbs.erase(it);
    }

    lua_pushboolean(L, true);
    return 1;
}

int cgzCLua::luaopen_sqlite3(lua_State *L)
{
    static const luaL_Reg sqlitelib[] = {
        {"new", sqlite_new},
        {"create", sqlite_create},
        {"drop", sqlite_drop},
        {"alter", sqlite_alter},
        {"truncate", sqlite_truncate},
        {"insert", sqlite_insert},
        {"select", sqlite_select},
        {"update", sqlite_update},
        {"delete", sqlite_delete},
        {"close", sqlite_close},
        {NULL, NULL}};

    luaL_newlib(L, sqlitelib);
    return 1;
}